package com.mdnote.structure.treeapplication;


import java.io.*;
import java.util.*;

/**
 * @author Rhythm-2019
 * @version 1.0
 * @date 2020/9/2
 * @description 赫夫曼编码实战 压缩文件
 */
public class HuffmanCodeZipFile {

    private static class Node implements Comparable<Node>{

        private Integer count;
        private Byte data;

        private Node left;
        private Node right;

        public Node(Byte letter, Integer count) {
            this.data = letter;
            this.count = count;
        }

        @Override
        public int compareTo(Node o) {
            return this.count - o.count;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "count=" + count +
                    ", letter=" + data +
                    ", left=" + left +
                    ", right=" + right +
                    '}';
        }
    }
    private static Map<Byte, String> huffmanCodeMap;

    private static int count = 0;

    public static void main(String[] args) {
        // 读取文件
        byte[] zipSourceData = readFile("D:\\test.png");
        System.out.println("压缩前文件大小：" + zipSourceData.length / 8 + "B");
        // 编码
        byte[] zipData = encode(zipSourceData);
        System.out.println("编码后文件大小为：" + zipData.length / 8 + "B");
        System.out.println(count);
        // 写回文件
        writeByObject("D:/test.zip", huffmanCodeMap, zipData);
        // 解压文件
        byte[] unZipSourceData = readByObject("D:/test.zip", huffmanCodeMap);
        // 反编码
        byte[] unZipData = decode(unZipSourceData);
        // 写文件
        writeFile(unZipData, "D:\\test.png");

    }

    /**
     * 写入文件
     * @param unZipData 解压后的数据
     * @param path 路径
     */
    private static void writeFile(byte[] unZipData, String path) {
        FileOutputStream fileOutputStream = null;
        try {
            fileOutputStream = new FileOutputStream(path);
            fileOutputStream.write(unZipData);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (fileOutputStream != null) {
                try {
                    fileOutputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

    }

    /**
     *
     * @param path
     * @param huffmanCodeMap
     * @return
     */
    @SuppressWarnings("uncheck")
    private static byte[] readByObject(String path, Map<Byte, String> huffmanCodeMap) {
        FileInputStream fileInputStream = null;
        ObjectInputStream objectInputStream = null;
        try {
            fileInputStream = new FileInputStream(path);
            objectInputStream = new ObjectInputStream(fileInputStream);
            byte[] bytes = (byte[]) objectInputStream.readObject();
            huffmanCodeMap = (Map<Byte, String>) objectInputStream.readObject();
            return bytes;
        } catch (IOException e) {
            e.printStackTrace();
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } finally {
            if (objectInputStream != null) {
                try {
                    objectInputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (fileInputStream != null) {
                try {
                    fileInputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return null;
    }

    /**
     * 写入文件
     * @param path 文件路径
     * @param huffmanCodeMap 编码表
     * @param bytes 数组
     */
    private static void writeByObject(String path, Map<Byte, String> huffmanCodeMap, byte[] bytes) {
        FileOutputStream fileOutputStream = null;
        ObjectOutputStream objectOutputStream = null;
        try {
            fileOutputStream = new FileOutputStream(path);
            objectOutputStream = new ObjectOutputStream(fileOutputStream);
            objectOutputStream.writeObject(bytes);
            objectOutputStream.writeObject(huffmanCodeMap);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (objectOutputStream != null) {
                try {
                    objectOutputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if (fileOutputStream != null) {
                try {
                    fileOutputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

    }

    private static byte[] encode(byte[] sourceArray) {
        // 构建哈夫曼树
        Node root = build(sourceArray);

        // 根据哈夫曼树构建哈夫曼编码表
        Map<Byte, String> huffmanCodeMap = generateCodeMap(root);
        HuffmanCodeZipFile.huffmanCodeMap = huffmanCodeMap;

        // 根据翻译源字符串
        String longBinaryString = toEncodeString(sourceArray, huffmanCodeMap);
        // 将这个二进制字符串转成二进制数组
        return toByteArray(longBinaryString);
    }

    private static Node build(byte[] sourceArray) {
        List<Node> nodeList = new ArrayList<Node>();

        Map<Byte, Integer> byteMap = new HashMap<>();
        for (byte b :sourceArray){
            Integer count = byteMap.get(b);
            if (count == null) {
                count = 1;
            } else {
                count += 1;
            }
            byteMap.put(b, count);
        }
        for (Map.Entry<Byte, Integer> entry : byteMap.entrySet()) {
            nodeList.add(new Node(entry.getKey(), entry.getValue()));
        }

        return buildHuffmanTree(nodeList);
    }

    /**
     * 读取文件内容
     * @param filePath 路径
     * @return 文件内容
     */
    private static byte[] readFile(String filePath) {

        FileInputStream fileInputStream = null;
        try {
            fileInputStream = new FileInputStream(filePath);
            byte[] bytes = new byte[fileInputStream.available()];
            fileInputStream.read(bytes);
            return bytes;
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (fileInputStream == null) {
                try {
                    fileInputStream.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return null;
    }

    /**
     * 反编码
     * @param array 编码数组
     * @return 反编码结果
     */
    private static byte[] decode(byte[] array) {
        // 把这个数组转回长字符创
        String decodeString = toDecodeString(array);
        // 翻转哈夫曼表的key和value
        HashMap<String, Byte> convertHuffmanCodeMap = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodeMap.entrySet()) {
            convertHuffmanCodeMap.put(entry.getValue(), entry.getKey());
        }
        // 按照哈夫曼表解码一下
        return decodeByHUffmanTree(decodeString, convertHuffmanCodeMap);
    }

    /**
     * 根据哈夫曼树反编码
     * @param decodeString 编码字符串
     * @param convertHuffmanCodeMap 哈夫曼编码表
     * @return 解码后的字符串
     */
    private static byte[] decodeByHUffmanTree(String decodeString, HashMap<String, Byte> convertHuffmanCodeMap) {

        int slowCursor = 0;
        int fastCursor = 1;
        List<Byte> byteList = new ArrayList<>();
        while (fastCursor <= decodeString.length()) {
            String temp = decodeString.substring(slowCursor, fastCursor);
            Byte Byte = convertHuffmanCodeMap.get(temp);
            if (Byte != null) {
                byteList.add(Byte);
                slowCursor = fastCursor;
            }
            fastCursor++;
        }
        byte[] bytes = new byte[byteList.size()];
        for (int i = 0; i < byteList.size(); i++) {
            bytes[i] = byteList.get(i);
        }
        return bytes;
    }

    /**\
     * 转成解码的字符串
     * @param array 数组
     * @return 字符串
     */
    private static String toDecodeString(byte[] array) {
        StringBuilder stringBuilder = new StringBuilder();
        // 有几种情况，byte为00110011，我们要补0
        for (int i = 0; i < array.length; i++) {
            int temp = array[i];
            String path = null;
            if (i + 1 == array.length) {
                // 最后一个元素
                path = Integer.toBinaryString(temp);
            } else {
                temp = array[i] | 256;;
                String binaryString = Integer.toBinaryString(temp);
                path = binaryString.substring(binaryString.length() - 8);
            }
            stringBuilder.append(path);
        }

        return stringBuilder.toString();
    }

    /**
     * 将二进制字符串转成数组
     * @param longBinaryString 二进制字符串转数组
     * @return 二进制数组
     */
    private static byte[] toByteArray(String longBinaryString) {
        int len = 0;
        if (longBinaryString.length() % 8 == 0) {
            len = longBinaryString.length() / 8;
        } else {
            len = longBinaryString.length() / 8 + 1;
        }
        byte[] bytes = new byte[len];
        int index = 0;
        for (int i = 0; i < longBinaryString.length(); i += 8) {
            if (i + 8 >= longBinaryString.length()) {
                // 最后一个8位不满8位
                bytes[index] = (byte) Integer.parseInt(longBinaryString.substring(i), 2);
            } else {
                bytes[index] = (byte) Integer.parseInt(longBinaryString.substring(i, i + 8), 2);
            }
            index ++;
        }

        return bytes;
    }

    /**
     * 根据哈夫曼表翻译字符串
     * @param sourceString 字符串
     * @param huffmanCodeMap 编码表
     * @return
     */
    private static String toEncodeString(byte[] sourceString, Map<Byte, String> huffmanCodeMap) {

        StringBuilder stringBuilder = new StringBuilder();
        for (byte b : sourceString) {
            stringBuilder.append(huffmanCodeMap.get(b));
        }
        return stringBuilder.toString();
    }

    /**
     * 生成哈夫曼表
     * @param root 根尖点
     * @return 哈夫曼表
     */
    private static Map<Byte, String> generateCodeMap(Node root) {
        Map<Byte, String> huffmanCodeMap = new HashMap<>();
        // 遍历哈夫曼树，把叶子节点加入哈夫曼表中
        if (root != null) {
            preOrder(root, huffmanCodeMap, "");
        }
//        for (Map.Entry<Byte, String> entry : huffmanCodeMap.entrySet()) {
//            System.out.println(entry.getKey() + " = " + entry.getValue());
//        }
        return huffmanCodeMap;
    }

    /**
     * 前序遍历哈夫曼树
     * @param node 节点
     * @param huffmanCodeMap 哈夫曼表
     * @param path 路径
     */
    private static void preOrder(Node node, Map<Byte, String> huffmanCodeMap, String path) {
        if (node.data != null) {
//            System.out.println(node);
            huffmanCodeMap.put(node.data, path);
            count++;
        }
        if (node.left != null) {
            preOrder(node.left, huffmanCodeMap, path + "0");
        }
        if (node.right != null) {
            preOrder(node.right, huffmanCodeMap, path + "1");
        }
    }



    /**
     * 构建哈夫曼树
     * @param nodeList 集合
     * @return 根节点
     */
    private static Node buildHuffmanTree(List<Node> nodeList) {

        // 排序
        Collections.sort(nodeList);
        while (nodeList.size() > 1) {
            // 拿出第一第二个元素
            Node leftNode = nodeList.get(0);
            Node rightNode = nodeList.get(1);
            // 构建新的节点
//            Node newNode = new Node(null, 0);
            Node newNode = new Node(null, leftNode.count + rightNode.count);
            newNode.left = leftNode;
            newNode.right = rightNode;
            // 删除、加入集合
            nodeList.remove(leftNode);
            nodeList.remove(rightNode);
            nodeList.add(newNode);
            // 重新排序
            Collections.sort(nodeList);
        }

        return nodeList.get(0);
    }
}
